{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "13f65a34",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.4.6.post5'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import sglang\n",
    "sglang.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c6503fb5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "Overlap scheduler is disabled because of using eagle speculative decoding.\n",
      "Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:02<00:00,  2.45s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:02<00:00,  2.45s/it]\n",
      "\n",
      "Loading pt checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n",
      "Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  1.74it/s]\n",
      "Loading pt checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  1.74it/s]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hot token id tensor([     0,      1,      2,  ..., 126437, 127196, 128009], device='cuda:0')\n"
     ]
    }
   ],
   "source": [
    "from sglang.srt.entrypoints.engine import Engine\n",
    "import torch\n",
    "llm = Engine(model_path=\"/models/Llama-3.2-1B-Instruct\",\n",
    "                speculative_algorithm = \"EAGLE3\",\n",
    "                speculative_draft_model_path = \"/models/sglang-EAGLE3-Llama-3.1-Instruct-8B\",\n",
    "                speculative_num_steps = 5,\n",
    "                speculative_eagle_topk = 8,\n",
    "                speculative_num_draft_tokens = 32,\n",
    "                max_running_requests = 8,\n",
    "                cuda_graph_max_bs = 8,\n",
    "                max_total_tokens = 4096,\n",
    "                disable_cuda_graph = True,\n",
    "                mem_fraction_static = 0.6,\n",
    "                dtype = torch.float16)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3759fb54",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.10/dist-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "配置加载成功。\n",
      "\n",
      "LlamaConfig 对象已创建.\n",
      "\n",
      "正在使用 torch.bfloat16 初始化 LlamaForCausalLM 模型...\n",
      "模型已成功使用随机权重进行初始化。\n",
      "模型参数量: 1,184,911,360\n",
      "词嵌入权重的数据类型: torch.float32\n",
      "模型的主要数据类型 (model.dtype): torch.float32\n",
      "\n",
      "模型将被保存到: /models/Dummy-Llama-3.2-8B\n",
      "正在尝试使用 safe_serialization=True 保存模型...\n",
      "模型和配置文件已成功保存到 /models/Dummy-Llama-3.2-8B/\n",
      "检查保存的文件：\n",
      " - .gitattributes\n",
      " - LICENSE.txt\n",
      " - README.md\n",
      " - generation_config.json\n",
      " - special_tokens_map.json\n",
      " - USE_POLICY.md\n",
      " - tokenizer.json\n",
      " - tokenizer_config.json\n",
      " - .cache\n",
      " - model.safetensors\n",
      " - config.json\n",
      "\n",
      "成功找到 /models/Dummy-Llama-3.2-8B/model.safetensors 文件。\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import torch\n",
    "from transformers import LlamaConfig, LlamaForCausalLM\n",
    "import os # 用于创建目录\n",
    "\n",
    "# --- 从上一个代码块复制模型初始化部分 ---\n",
    "\n",
    "# 0. 您的配置 (作为多行字符串)\n",
    "config_json_string = \"\"\"\n",
    "{\n",
    "  \"architectures\": [\n",
    "    \"LlamaForCausalLM\"\n",
    "  ],\n",
    "  \"attention_bias\": false,\n",
    "  \"attention_dropout\": 0.0,\n",
    "  \"bos_token_id\": 128000,\n",
    "  \"eos_token_id\": 128009,\n",
    "  \"hidden_act\": \"silu\",\n",
    "  \"hidden_size\": 4096,\n",
    "  \"initializer_range\": 0.02,\n",
    "  \"intermediate_size\": 2048,\n",
    "  \"max_position_embeddings\": 8192,\n",
    "  \"model_type\": \"llama\",\n",
    "  \"num_attention_heads\": 32,\n",
    "  \"num_hidden_layers\": 2,\n",
    "  \"num_key_value_heads\": 8,\n",
    "  \"pretraining_tp\": 1,\n",
    "  \"rms_norm_eps\": 1e-05,\n",
    "  \"rope_scaling\": null,\n",
    "  \"rope_theta\": 500000.0,\n",
    "  \"tie_word_embeddings\": false,\n",
    "  \"torch_dtype\": \"bfloat16\",\n",
    "  \"transformers_version\": \"4.40.0.dev0\",\n",
    "  \"use_cache\": true,\n",
    "  \"vocab_size\": 128256\n",
    "}\n",
    "\"\"\"\n",
    "\n",
    "# 1. 将 JSON 配置字符串解析为 Python 字典\n",
    "config_dict = json.loads(config_json_string)\n",
    "print(\"配置加载成功。\")\n",
    "\n",
    "# 2. 创建 LlamaConfig 对象\n",
    "llama_custom_config = LlamaConfig(**config_dict)\n",
    "print(f\"\\nLlamaConfig 对象已创建.\")\n",
    "\n",
    "# 3. 使用配置初始化 LlamaForCausalLM 模型\n",
    "print(f\"\\n正在使用 {llama_custom_config.torch_dtype} 初始化 LlamaForCausalLM 模型...\")\n",
    "model = LlamaForCausalLM(config=llama_custom_config)\n",
    "print(\"模型已成功使用随机权重进行初始化。\")\n",
    "\n",
    "# 验证初始化模型的一些属性\n",
    "print(f\"模型参数量: {model.num_parameters():,}\")\n",
    "if hasattr(model, 'model') and hasattr(model.model, 'embed_tokens'):\n",
    "    print(f\"词嵌入权重的数据类型: {model.model.embed_tokens.weight.dtype}\")\n",
    "print(f\"模型的主要数据类型 (model.dtype): {model.dtype}\")\n",
    "\n",
    "# --- 新增：保存模型为 safetensors 格式 ---\n",
    "\n",
    "# 5. 定义保存路径\n",
    "save_directory = \"/models/Dummy-Llama-3.2-8B/\"\n",
    "if not os.path.exists(save_directory):\n",
    "    os.makedirs(save_directory)\n",
    "print(f\"\\n模型将被保存到: {os.path.abspath(save_directory)}\")\n",
    "\n",
    "# 6. 保存模型\n",
    "# `save_pretrained` 方法会自动保存配置文件 (config.json) 和模型权重。\n",
    "# 通过设置 `safe_serialization=True`，模型权重将被保存为 `model.safetensors`。\n",
    "# 如果transformers版本较旧，可能默认已经是safetensors或需要显式安装safetensors库。\n",
    "# 确保你的transformers版本支持 `safe_serialization` 参数。\n",
    "# 对于 \"4.40.0.dev0\"，这应该是支持的。\n",
    "print(\"正在尝试使用 safe_serialization=True 保存模型...\")\n",
    "model = model.cuda().to(torch.bfloat16)\n",
    "model.save_pretrained(save_directory, safe_serialization=True)\n",
    "print(f\"模型和配置文件已成功保存到 {save_directory}\")\n",
    "print(\"检查保存的文件：\")\n",
    "for f_name in os.listdir(save_directory):\n",
    "    print(f\" - {f_name}\")\n",
    "\n",
    "# 验证是否生成了 .safetensors 文件\n",
    "safetensors_file_path = os.path.join(save_directory, \"model.safetensors\")\n",
    "if os.path.exists(safetensors_file_path):\n",
    "    print(f\"\\n成功找到 {safetensors_file_path} 文件。\")\n",
    "else:\n",
    "    # 如果 safe_serialization=True 失败或不被支持，它可能会回退到 .bin\n",
    "    bin_file_path = os.path.join(save_directory, \"pytorch_model.bin\")\n",
    "    if os.path.exists(bin_file_path):\n",
    "        print(f\"\\n警告: 未找到 model.safetensors。但找到了 {bin_file_path}。\")\n",
    "        print(\"请确保您的 Transformers 和 safetensors 库是最新的，以支持 safe_serialization=True。\")\n",
    "        print(\"如果您的环境不支持，模型可能已保存为 pytorch_model.bin。\")\n",
    "    else:\n",
    "        print(\"\\n错误: 未找到 model.safetensors 或 pytorch_model.bin 文件。保存可能未成功。\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1014848b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import LlamaForCausalLM\n",
    "\n",
    "model = LlamaForCausalLM.from_pretrained(\"/models/Dummy-Llama-3.2-8B\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
